#include "pe_parser.h"

#include "macro.h"

#include <array>
#include <stdexcept>

namespace analyzer::windows
{

namespace
{

template <typename data>
bool is_null(const data&);

template <uint8_t bitness>
bool is_null(const image_thunk_data<bitness>& data)
{
    return data.u1.function == 0;
}

template <>
bool is_null(const image_import_descriptor& descriptor)
{
    return descriptor.dummy.characteristics == 0;
}

uint32_t read_position_in_section(const image_section_header& section, const uint32_t rva)
{
    return rva - section.virtual_address + section.pointer_to_raw_data;
}

template <uint8_t bitness>
std::vector<image_thunk_data<bitness>> read_thunks(std::ifstream& file, const image_section_header& section, const uint32_t rva)
{
    const auto function_names_position = read_position_in_section(section, rva);
    file.seekg(function_names_position);
    std::vector<image_thunk_data<bitness>> thunks;
    image_thunk_data<bitness> thunk_data{};

    while (file.good())
    {
        file.read(std::launder(reinterpret_cast<char*>(&thunk_data)), sizeof(thunk_data));
        if (is_null(thunk_data))
            break;

        thunks.push_back(thunk_data);
    }

    return thunks;
}

template <uint8_t bitness>
image_import_by_name read_import_by_name(std::ifstream& file, const image_section_header& section, const uint32_t rva)
{
    const auto function_name_position = read_position_in_section(section, rva);
    uint16_t hint;
    file.seekg(function_name_position);
    file.read(std::launder(reinterpret_cast<char*>(&hint)), sizeof(hint));

    char c{};
    std::string name;
    file.read(&c, sizeof(c));
    while (c)
    {
        name += c;
        file.read(&c, sizeof(c));
    }

    return image_import_by_name{.hint = hint, .name = std::move(name)};
}

template <uint8_t bitness>
std::vector<std::string> read_import_functions(std::ifstream& file, const image_section_header& section, const uint32_t rva)
{
    std::vector<std::string> function_names;
    for (const auto& thunk_data : read_thunks<bitness>(file, section, rva))
    {
        auto function_info = read_import_by_name<bitness>(file, section, thunk_data.u1.function);
        function_names.push_back(std::move(function_info.name));
    }

    return function_names;
}
} // anonymous namespace

pe_parser::pe_parser(std::filesystem::path path) : path(std::move(path))
{
}

std::optional<file_structure> pe_parser::parse()
{
    file = std::ifstream(path, std::ios::binary);
    if (!file.is_open())
        throw std::runtime_error("can't open file");

    if (!contains_magic_bytes())
        return std::nullopt;

    read_image_header(read_image_header_position());

    uint16_t number_of_sections{};
    uint32_t section_alignment{};
    uint32_t import_section_virtual_address{};

    switch (bitness)
    {
    case x86:
    {
        const auto& header32 = std::get<image_nt_header<x86>>(image_header);
        number_of_sections = header32.file_header.number_of_sections;
        section_alignment = header32.optional_header.section_alignment;
        import_section_virtual_address = header32.optional_header.data_directory[image_directory_entry_import].virtual_address;
        break;
    }
    case x64:
    {
        const auto& header64 = std::get<image_nt_header<x64>>(image_header);
        number_of_sections = header64.file_header.number_of_sections;
        section_alignment = header64.optional_header.section_alignment;
        import_section_virtual_address = header64.optional_header.data_directory[image_directory_entry_import].virtual_address;
        break;
    }
    }

    read_section_headers(number_of_sections);
    import_section_index = read_import_index(section_alignment, import_section_virtual_address);

    const auto& section = section_headers[import_section_index];

    const auto import_section_position = read_position_in_section(section, import_section_virtual_address);

    read_import_descriptors(import_section_position);

    return file_structure{
        .format = file_structure::pe,
        .bitness = bitness,
        .import_modules = read_import_modules()};
}

bool pe_parser::contains_magic_bytes()
{
    std::array<char, 2> magic;
    file.read(magic.data(), magic.size() * sizeof(decltype(magic)::value_type));
    return magic[0] == 'M' && magic[1] == 'Z';
}

void pe_parser::read_bitness()
{
    uint16_t magic;
    file.read(reinterpret_cast<char*>(&magic), sizeof(magic));
    file.seekg(-sizeof(magic), std::ios_base::cur);

    switch (magic)
    {
    case image_nt_optional_hdr32_magic:
        bitness = x86;
        break;
    case image_nt_optional_hdr64_magic:
        bitness = x64;
        break;
    default:
        throw std::runtime_error("unknown bitness");
    }
}

uint16_t pe_parser::read_image_header_position()
{
    file.seekg(0x3C);
    uint16_t header_position{};
    file.read(reinterpret_cast<char*>(&header_position), sizeof(header_position));

    return header_position;
}

void pe_parser::read_image_header(const uint16_t position)
{
    file.seekg(position);

    uint32_t signature{};
    file.read(reinterpret_cast<char*>(&signature), sizeof(signature));
    if (signature != 0x4550)
        throw std::runtime_error("can't read signature of image header");

    image_file_header file_header;
    file.read(reinterpret_cast<char*>(&file_header), sizeof(file_header));
    read_bitness();

    switch (bitness)
    {
    case x86:
    {
        image_optional_header32 optional_header;
        file.read(reinterpret_cast<char*>(&optional_header), sizeof(optional_header));
        image_header = image_nt_header<x86>{
            .signature = signature,
            .file_header = file_header,
            .optional_header = optional_header};
        break;
    }
    case x64:
    {
        image_optional_header64 optional_header;
        file.read(reinterpret_cast<char*>(&optional_header), sizeof(optional_header));
        image_header = image_nt_header<x64>{
            .signature = signature,
            .file_header = file_header,
            .optional_header = optional_header};
        break;
    }
    }
}

void pe_parser::read_section_headers(const uint16_t number_of_sections)
{
    section_headers.resize(number_of_sections);
    for (auto& section_header : section_headers)
        file.read(reinterpret_cast<char*>(&section_header), sizeof(section_header));
}

void pe_parser::read_import_descriptors(const uint32_t import_section_position)
{
    image_import_descriptor descriptor{};
    file.seekg(import_section_position);

    while (file.good())
    {
        file.read(reinterpret_cast<char*>(&descriptor), sizeof(descriptor));
        if (is_null(descriptor))
            break;

        import_descriptors.push_back(descriptor);
    }
}

std::string pe_parser::read_module_name(const image_section_header& section, const uint32_t name_address)
{
    const auto name_position = read_position_in_section(section, name_address);
    file.seekg(name_position);
    std::string name;
    char c{};

    do
    {
        file.read(&c, sizeof(c));
        name.push_back(c);
    } while (c && file.good());

    return name;
}

std::size_t pe_parser::read_import_index(const uint32_t alignment, const uint32_t rva)
{
    const auto align = [alignment](const auto value)
    {
        if (const auto remainder = value % alignment; remainder)
            return value + alignment - remainder;

        return value;
    };

    for (std::size_t i = 0; i < section_headers.size(); ++i)
    {
        const auto& section = section_headers[i];
        if (rva >= section.virtual_address && rva < section.virtual_address + align(section.misc.virtual_size))
            return i;
    }

    throw std::runtime_error("can't find sections index");
}

std::vector<import_module> pe_parser::read_import_modules()
{
    const auto& section = section_headers[import_section_index];
    std::vector<import_module> import_modules;

    for (const auto& descriptor : import_descriptors)
    {
        auto module_name = read_module_name(section, descriptor.name);
        auto functions = APPLY_DEPENDING_ON_BITNESS(read_import_functions, bitness)(file, section, descriptor.dummy.original_first_thunk);
        import_modules.push_back({std::move(module_name), std::move(functions)});
    }

    return import_modules;
}

} // namespace analyzer::windows
